import pandas as pd
import numpy as np
import talib
import time
import psycopg2
import sys


conn = psycopg2.connect(database="finance",
                        user="postgres", password="123456",
                        host="127.0.0.1", port="5432")
cur = conn.cursor()  # 创建指针对象

stock_rsi = []
sql = ""

cur.execute("SELECT code, sname FROM stock_market.stock")
datas = cur.fetchall()
for item in datas:
    code = item[0]
    sname = item[1]
    if "ST" in sname:
        print("不计算 %s", sname)
        continue
    if "退" in sname:
        print("不计算 %s", sname)
        continue
    try:
        select_sql = '''SELECT id,date,open,close,hight,low FROM stock_datas.stock_qfq_{} ORDER BY date DESC LIMIT {};'''
        sql = select_sql.format(code, 300)
        # print(sql)
        cur.execute(sql)
        df = pd.DataFrame(cur.fetchall())
        if df.values.size <= 10:
            print("======", code, sname, df.values.size)
            continue
        df.columns= ["id", "date", "open", "close", "high", "low"]
        new_df = df.sort_values("date", ascending=True)
        col = new_df.iloc[:, 3]
        arr = col.values
        rsi = talib.RSI(arr, timeperiod=14)
        if len(rsi) < 100:
            continue
        rsi20 = rsi[-2:]
        for item in rsi20:
            if item <= 25.0:
                stock_rsi.append(code)
                print("insert", code, rsi20)
                break
        # new_df.insert(6, "rsi", rsi)
        # print(new_df.tail(30))
    except Exception as e:
        print(sql)
        cur.execute("ROLLBACK")
        print('数据库commit：{}'.format(e))
        sys.exit(1)
    time.sleep(0.001)
print("============================================")
print(stock_rsi)
print("============================================")